Skip to content

Add MLX dispatch for CholeskySolve - #2377

Merged
jessegrabowski merged 2 commits into
pymc-devs:mainfrom
jessegrabowski:mlx-cho-solve
Aug 26, 2026
Merged

Add MLX dispatch for CholeskySolve#2377
jessegrabowski merged 2 commits into
pymc-devs:mainfrom
jessegrabowski:mlx-cho-solve

Conversation

@jessegrabowski

Copy link
Copy Markdown
Member

MLX has no cho_solve, so this runs two triangular solves against the Cholesky factor and its transpose.

Every op in the dispatch, including the transpose, is pinned to the CPU stream. mlx.linalg is CPU-only, and a transpose left on the default stream aborts on float64 input.

@jessegrabowski
jessegrabowski merged commit 73e5794 into pymc-devs:main Aug 26, 2026
67 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants